#!/usr/bin/env python3
# coding:utf-8
import os
import sys

from vision.msg import darknetMes
import cv2
from GetCvPicture import PicGetter
import rospy
import numpy as np
import time
sys.path.insert(0, "/home/ucar/ucar_ws/src/darknet_pkg/src/darknet_yolov4_tiny")
from darknet_yolov4_tiny.Detect import Detect


class Detector:
    def __init__(self):
        rospy.init_node('Darknet_publisher')

        # Darknet目标检测实体对象
        self.detect = Detect()
        self.base_conf = 0.6
        self.cap = PicGetter()
        self.pub = None
        self.openFlag = False

    def Predict(self, img):
        results, plot = self.detect.predict_image(img, self.base_conf)

        mes = darknetMes()

        for res in results:
            mes.name.append(str(res[0]))
            mes.conf.append(float(res[1]))
            mes.x.append(float(res[2][0]))
            mes.y.append(float(res[2][1]))
            mes.w.append(float(res[2][2]))
            mes.h.append(float(res[2][3]))

        return mes

    def start(self):
        self.cap.start()

        self.pub = rospy.Publisher('DarknetMes', darknetMes, queue_size=2)

    def main(self):
        while not rospy.is_shutdown():
            self.openFlag = rospy.get_param("OpenDarknet", False)
            if not self.openFlag:
                time.sleep(0.1)
                continue
            img = self.cap.Get()

            if img is None:
                time.sleep(0.1)
                continue

            start = time.time()
            mes = self.Predict(img)

            self.pub.publish(mes)

            # print(f"costtime:{time.time() - start}")
            # print(f"fps:{1/(time.time() - start)}")


if __name__ == '__main__':
   detector = Detector()

   detector.start()

   detector.main()
